"""
财务报表和估值指标相关的Celery任务
"""
import logging
import datetime
from celery import group, chain
from celery.utils.log import get_task_logger
from sqlalchemy import func

from modules.tasks.celery_app import app
from modules.models.base import get_db
from modules.models.stock import SecurityInfo
from modules.data_collector import collector
from modules.tasks.news_tasks import collect_stock_notice

# 配置任务日志
logger = get_task_logger(__name__)


@app.task(bind=True, queue='stock_data', max_retries=3, retry_backoff=60)
def collect_financial_statement(self, stock_code, report_type="balance", period="年报"):
    """
    采集财务报表数据
    
    Args:
        stock_code: 股票代码
        report_type: 报表类型，可选值：balance(资产负债表)/income(利润表)/cash_flow(现金流量表)
        period: 报表周期，可选值：年报/一季报/中报/三季报
    """
    logger.info(f"开始采集股票 {stock_code} 的{report_type}财务报表数据，周期: {period}")
    try:
        count = collector.collect_market_data(
            "financial_statement", 
            code=stock_code,
            report_type=report_type,
            period=period
        )
        logger.info(f"成功采集{count}条股票 {stock_code} 的{report_type}财务报表数据")
        return {
            "status": "success",
            "stock_code": stock_code,
            "report_type": report_type,
            "period": period,
            "count": count
        }
    except Exception as e:
        logger.error(f"采集股票 {stock_code} 的{report_type}财务报表数据失败: {str(e)}")
        raise self.retry(exc=e)


@app.task(bind=True, queue='stock_data', max_retries=3, retry_backoff=60)
def collect_valuation_indicators(self, stock_code):
    """
    采集估值指标数据
    
    Args:
        stock_code: 股票代码
    """
    logger.info(f"开始采集股票 {stock_code} 的估值指标数据")
    try:
        count = collector.collect_market_data(
            "valuation", 
            code=stock_code
        )
        logger.info(f"成功采集{count}条股票 {stock_code} 的估值指标数据")
        return {
            "status": "success",
            "stock_code": stock_code,
            "count": count
        }
    except Exception as e:
        logger.error(f"采集股票 {stock_code} 的估值指标数据失败: {str(e)}")
        raise self.retry(exc=e)


@app.task(bind=True, queue='stock_data')
def schedule_financial_data_collection(self, stock_list=None, max_stocks=50):
    """
    调度财务数据采集任务
    
    Args:
        stock_list: 股票代码列表，默认为None (自动选择)
        max_stocks: 最大采集股票数量
    """
    logger.info("开始调度财务数据采集任务")
    
    try:
        # 如果未提供股票列表，则从数据库获取
        db = next(get_db())
        
        if stock_list is None:
            # 获取市值排名靠前的股票
            stocks = db.query(SecurityInfo).filter(
                SecurityInfo.type == 'stock',
                SecurityInfo.is_active == 1
            ).order_by(
                # 按ID排序，假设ID越小市值越大（实际可能需要按其他字段排序）
                SecurityInfo.id
            ).limit(max_stocks).all()
            
            stock_list = [stock.code for stock in stocks]
        else:
            # 限制股票数量
            stock_list = stock_list[:max_stocks]
        
        db.close()
        
        logger.info(f"准备采集 {len(stock_list)} 只股票的财务数据")
        
        # 创建任务链
        tasks = []
        for stock_code in stock_list:
            # 对每只股票，采集三种财务报表和估值指标
            subtasks = chain(
                collect_financial_statement.s(stock_code, "balance", "年报"),
                collect_financial_statement.s(stock_code, "income", "年报"),
                collect_financial_statement.s(stock_code, "cash_flow", "年报"),
                collect_valuation_indicators.s(stock_code)
            )
            tasks.append(subtasks)
        
        # 分批执行任务
        batch_size = 10
        for i in range(0, len(tasks), batch_size):
            batch_tasks = tasks[i:i+batch_size]
            for task in batch_tasks:
                task.apply_async()
            logger.info(f"已调度 {i+len(batch_tasks)}/{len(tasks)} 只股票的财务数据采集任务")
        
        # 同时采集公告
        for stock_code in stock_list:
            collect_stock_notice.delay(stock_code)
        
        return {
            "status": "success",
            "message": f"已调度 {len(stock_list)} 只股票的财务数据采集任务"
        }
    except Exception as e:
        logger.error(f"调度财务数据采集任务失败: {str(e)}")
        raise


@app.task(bind=True, queue='stock_data')
def collect_all_financial_reports(self, stock_code, periods=None):
    """
    采集股票的所有类型财务报表
    
    Args:
        stock_code: 股票代码
        periods: 报表周期列表，默认为["年报", "中报"]
    """
    logger.info(f"开始采集股票 {stock_code} 的所有财务报表")
    
    try:
        if periods is None:
            periods = ["年报", "中报"]
        
        report_types = ["balance", "income", "cash_flow"]
        
        # 创建任务组
        tasks = []
        for period in periods:
            for report_type in report_types:
                tasks.append(collect_financial_statement.s(stock_code, report_type, period))
        
        # 添加估值指标采集任务
        tasks.append(collect_valuation_indicators.s(stock_code))
        
        # 执行任务组
        group(tasks).apply_async()
        
        logger.info(f"已调度股票 {stock_code} 的 {len(tasks)} 个财务数据采集任务")
        return {
            "status": "success",
            "message": f"已调度股票 {stock_code} 的 {len(tasks)} 个财务数据采集任务",
            "stock_code": stock_code
        }
    except Exception as e:
        logger.error(f"采集股票 {stock_code} 的财务报表失败: {str(e)}")
        raise


# 测试函数
def test_financial_tasks():
    """测试财务数据任务函数"""
    # 采集单只股票的财务报表
    db = next(get_db())
    stock = db.query(SecurityInfo).filter(
        SecurityInfo.type == 'stock'
    ).first()
    db.close()
    
    if stock:
        # 测试采集资产负债表
        statement_result = collect_financial_statement.delay(stock.code, "balance", "年报")
        print(f"采集股票 {stock.code} 资产负债表任务已提交，任务ID: {statement_result.id}")
        
        # 等待任务完成
        result = statement_result.get(timeout=120)
        print(f"采集股票资产负债表结果: {result}")
        
        # 测试采集估值指标
        valuation_result = collect_valuation_indicators.delay(stock.code)
        print(f"采集股票 {stock.code} 估值指标任务已提交，任务ID: {valuation_result.id}")
        
        # 等待任务完成
        result = valuation_result.get(timeout=120)
        print(f"采集估值指标结果: {result}")
        
        # 测试采集所有财务报表
        all_reports_result = collect_all_financial_reports.delay(stock.code)
        print(f"采集股票 {stock.code} 所有财务报表任务已提交，任务ID: {all_reports_result.id}")
    else:
        print("没有找到股票数据，无法测试采集任务")


if __name__ == "__main__":
    # 配置日志
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
    
    # 执行测试
    test_financial_tasks()